#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
用户反馈学习模块
实现基于用户反馈的增量训练机制
"""

import os
import json
import sqlite3
from typing import List, Tuple
from datetime import datetime
from transformers import BertTokenizer, BertForSequenceClassification, AdamW
from torch.utils.data import Dataset, DataLoader
import torch


class FeedbackDataset(Dataset):
    """用户反馈数据集"""
    
    def __init__(self, texts: List[str], labels: List[int], tokenizer, max_length: int = 128):
        self.texts = texts
        self.labels = labels
        self.tokenizer = tokenizer
        self.max_length = max_length
    
    def __len__(self):
        return len(self.texts)
    
    def __getitem__(self, idx):
        text = str(self.texts[idx])
        label = self.labels[idx]
        
        encoding = self.tokenizer.encode_plus(
            text,
            add_special_tokens=True,
            max_length=self.max_length,
            padding='max_length',
            truncation=True,
            return_tensors='pt'
        )
        
        return {
            'input_ids': encoding['input_ids'].flatten(),
            'attention_mask': encoding['attention_mask'].flatten(),
            'label': torch.tensor(label, dtype=torch.long)
        }


class FeedbackLearner:
    """用户反馈学习器"""
    
    def __init__(self, model_path: str = '../model_cache/bert_transaction_classifier'):
        self.model_path = model_path
        self.feedback_db_path = '../model_cache/feedback.db'
        self.labels = ['餐饮', '交通', '购物', '娱乐', '医疗', '教育', '住房', '其他']
        self.label_map = {label: i for i, label in enumerate(self.labels)}
        
        # 确保model_cache目录存在
        os.makedirs('./model_cache', exist_ok=True)
        self.init_feedback_db()
        
    def init_feedback_db(self):
        """初始化反馈数据库"""
        conn = sqlite3.connect(self.feedback_db_path)
        cursor = conn.cursor()
        
        # 创建反馈记录表
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS feedback_records (
                id INTEGER PRIMARY KEY,
                text TEXT NOT NULL,
                predicted_category TEXT,
                user_corrected_category TEXT,
                timestamp DATETIME NOT NULL,
                is_trained BOOLEAN DEFAULT FALSE
            )
        """)
        
        conn.commit()
        conn.close()
        
    def add_feedback(self, text: str, predicted_category: str, user_corrected_category: str):
        """添加用户反馈记录"""
        conn = sqlite3.connect(self.feedback_db_path)
        cursor = conn.cursor()
        
        cursor.execute("""
            INSERT INTO feedback_records 
            (text, predicted_category, user_corrected_category, timestamp)
            VALUES (?, ?, ?, ?)
        """, (text, predicted_category, user_corrected_category, datetime.now()))
        
        conn.commit()
        conn.close()
        
    def get_untrained_feedback(self) -> List[Tuple]:
        """获取未训练的反馈记录"""
        conn = sqlite3.connect(self.feedback_db_path)
        cursor = conn.cursor()
        
        cursor.execute("""
            SELECT text, user_corrected_category 
            FROM feedback_records 
            WHERE is_trained = FALSE
        """)
        
        results = cursor.fetchall()
        conn.close()
        return results
        
    def mark_as_trained(self, texts: List[str]):
        """标记反馈记录为已训练"""
        conn = sqlite3.connect(self.feedback_db_path)
        cursor = conn.cursor()
        
        for text in texts:
            cursor.execute("""
                UPDATE feedback_records 
                SET is_trained = TRUE 
                WHERE text = ?
            """, (text,))
        
        conn.commit()
        conn.close()
        
    def incremental_train(self, epochs: int = 1, batch_size: int = 2, learning_rate: float = 2e-5):
        """增量训练模型"""
        # 获取未训练的反馈数据
        feedback_data = self.get_untrained_feedback()
        
        if not feedback_data:
            print("没有新的反馈数据需要训练")
            return
        
        # 准备训练数据
        texts = [item[0] for item in feedback_data]
        label_names = [item[1] for item in feedback_data]
        label_ids = [self.label_map[name] for name in label_names]
        
        # 加载现有模型和tokenizer
        try:
            tokenizer = BertTokenizer.from_pretrained(self.model_path, local_files_only=True)
            model = BertForSequenceClassification.from_pretrained(self.model_path, num_labels=len(self.labels), local_files_only=True)
        except Exception as e:
            print(f"加载模型时出错: {e}")
            # 尝试使用绝对路径
            abs_model_path = os.path.abspath(self.model_path)
            print(f"尝试使用绝对路径: {abs_model_path}")
            tokenizer = BertTokenizer.from_pretrained(abs_model_path, local_files_only=True)
            model = BertForSequenceClassification.from_pretrained(abs_model_path, num_labels=len(self.labels), local_files_only=True)
        
        # 创建数据集和数据加载器
        dataset = FeedbackDataset(texts, label_ids, tokenizer)
        dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
        
        # 设置优化器
        optimizer = AdamW(model.parameters(), lr=learning_rate)
        
        # 训练模型
        model.train()
        for epoch in range(epochs):
            for batch in dataloader:
                optimizer.zero_grad()
                input_ids = batch['input_ids']
                attention_mask = batch['attention_mask']
                labels = batch['label']
                
                outputs = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels)
                loss = outputs.loss
                loss.backward()
                optimizer.step()
                
                print(f"Epoch: {epoch}, Loss: {loss.item()}")
        
        # 保存模型
        model.save_pretrained(self.model_path)
        tokenizer.save_pretrained(self.model_path)
        print(f"模型已增量训练并保存到 {self.model_path}")
        
        # 标记反馈数据为已训练
        self.mark_as_trained(texts)


def main():
    """主函数，用于测试反馈学习功能"""
    learner = FeedbackLearner()
    
    # 添加一些测试反馈数据
    learner.add_feedback("在麦当劳用餐", "购物", "餐饮")
    learner.add_feedback("乘坐公交车", "交通", "交通")
    learner.add_feedback("购买书籍", "购物", "教育")
    
    # 执行增量训练
    learner.incremental_train(epochs=2)


if __name__ == "__main__":
    main()